﻿/*
** Xin YUAN, 2023, BSD (2)
*/

////////////////////////////////////////////////////////////////////////////////

#include <sys/socket.h>

#include <netdb.h>
#include <arpa/inet.h>

#include <fcntl.h>

////////////////////////////////////////////////////////////////////////////////

//_x_socket_addr

class _x_socket_addr
{
public:
	_x_socket_addr() noexcept : m_pInfo(NULL)
	{
	}
	_x_socket_addr(const _x_socket_addr&) = delete;
	_x_socket_addr(_x_socket_addr&& src) noexcept : m_pInfo(src.m_pInfo)
	{
		src.m_pInfo = NULL;
	}
	~_x_socket_addr() noexcept
	{
		Free();
	}

	_x_socket_addr& operator=(const _x_socket_addr&) = delete;
	_x_socket_addr& operator=(_x_socket_addr&& src) noexcept
	{
		if ( this != &src ) {
			if ( m_pInfo == src.m_pInfo ) {
				assert( m_pInfo == NULL );
			}
			else {
				Free();
				m_pInfo = src.m_pInfo;
				src.m_pInfo = NULL;
			}
		}
		return *this;
	}

	void Free() noexcept
	{
		if ( m_pInfo != NULL ) {
			::freeaddrinfo(m_pInfo);
			m_pInfo = NULL;
		}
	}

	bool IsNull() const noexcept
	{
		return m_pInfo == NULL;
	}
	const struct addrinfo* GetAddrInfo() const noexcept
	{
		return m_pInfo;
	}
	struct addrinfo* GetAddrInfo() noexcept
	{
		return m_pInfo;
	}

	bool Fetch(const char* node, const char* service) noexcept
	{
		assert( IsNull() );
		assert( node != NULL );
		struct addrinfo hints;
		::memset(&hints, 0, sizeof(struct addrinfo));
		return ::getaddrinfo(node, service, &hints, &m_pInfo)
					== 0;  //false, return value is error code
	}

protected:
	struct addrinfo* m_pInfo;
};

//_x_socket_helper

class _x_socket_helper
{
public:
	//hostlen: NI_MAXHOST  servlen: NI_MAXSERV
	static bool GetHostName(const struct sockaddr* addr, socklen_t addrlen,
					char* host, socklen_t hostlen,
					char* serv, socklen_t servlen,
					int flags = 0) noexcept
	{
		assert( addr != NULL && addrlen > 0 );
		return ::getnameinfo(addr, addrlen, host, hostlen, serv, servlen, flags)
					== 0;  //false, return value is error code
	}
	//size: INET_ADDRSTRLEN or INET6_ADDRSTRLEN
	static bool GetIPName(const struct sockaddr* addr, socklen_t addrlen,
					char* dest, socklen_t size) noexcept
	{
		assert( addr != NULL && addrlen > 0 );
		int af = ((const struct sockaddr_in*)addr)->sin_family;
		return ::inet_ntop(af,
					af == AF_INET6 ? (const void*)&(((const struct sockaddr_in6*)addr)->sin6_addr)
							: (const void*)&(((const struct sockaddr_in*)addr)->sin_addr),
					dest, size) != NULL;  //NULL, errno
	}
	static u_int16_t GetPort(const struct sockaddr* addr, socklen_t addrlen) noexcept
	{
		assert( addr != NULL && addrlen > 0 );
		return ::ntohs(((const struct sockaddr_in*)addr)->sin_port);
	}
	static void SetPort(u_int16_t port, struct sockaddr* addr, socklen_t addrlen) noexcept
	{
		assert( addr != NULL && addrlen > 0 );
		((struct sockaddr_in*)addr)->sin_port = ::htons(port);
	}
	static bool GetLocalHostName(char* name, size_t len) noexcept
	{
		assert( name != NULL && len > 0 );
		return ::gethostname(name, len) == 0;  //-1, errno
	}
	static bool FillSockAddr(struct sockaddr_storage& addr, socklen_t& addrlen,
					int af, u_int16_t port, const char* ip = NULL) noexcept
	{
		addrlen = sizeof(struct sockaddr_storage);
		::memset(&addr, 0, addrlen);

		((struct sockaddr_in*)&addr)->sin_family = af;
		if ( ip == NULL ) {
			if ( af == AF_INET6 )
				((struct sockaddr_in6*)&addr)->sin6_addr = IN6ADDR_ANY_INIT;
			else
				((struct sockaddr_in*)&addr)->sin_addr.s_addr = ::htonl(INADDR_ANY);
		}
		else {
			if ( ::inet_pton(af, ip,
					af == AF_INET6 ? (void*)&(((struct sockaddr_in6*)&addr)->sin6_addr)
							: (void*)&(((struct sockaddr_in*)&addr)->sin_addr)
					) != 1 )  //0: invalid input string  -1: errno
				return false;
		}

		addrlen = af == AF_INET6 ? sizeof(struct sockaddr_in6) : sizeof(struct sockaddr_in);
		SetPort(port, (struct sockaddr*)&addr, addrlen);
		return true;
	}
};

//_x_nb_socket

class _x_nb_socket
{
public:
	explicit _x_nb_socket(int fd = -1) noexcept : m_fd(fd)
	{
	}
	_x_nb_socket(const _x_nb_socket&) = delete;
	_x_nb_socket(_x_nb_socket&& src) noexcept : m_fd(src.m_fd)
	{
		src.m_fd = -1;
	}
	~_x_nb_socket() noexcept
	{
		Close();
	}

	_x_nb_socket& operator=(const _x_nb_socket&) = delete;
	_x_nb_socket& operator=(_x_nb_socket&& src) noexcept
	{
		if ( this != &src ) {
			if ( m_fd == src.m_fd ) {
				assert( m_fd == -1 );
			}
			else {
				Close();
				m_fd = src.m_fd;
				src.m_fd = -1;
			}
		}
		return *this;
	}

	void Close() noexcept
	{
		if ( m_fd != -1 ) {
			int ret = ::close(m_fd);
			(void)ret;
			assert( ret == 0 );  //-1, errno
			m_fd = -1;
		}
	}

	void GraceClose() noexcept
	{
		if ( IsValid() )
			Disable(SHUT_RDWR);  //no check
		Close();
	}

	operator int() const noexcept
	{
		return m_fd;
	}
	int GetHandle() const noexcept
	{
		return m_fd;
	}
	bool IsValid() const noexcept
	{
		return m_fd != -1;
	}

	void Attach(int fd) noexcept
	{
		if ( fd != m_fd ) {
			Close();
			m_fd = fd;
		}
	}
	int Detach() noexcept
	{
		int fd = m_fd;
		m_fd = -1;
		return fd;
	}

	//domain: AF_*
	//type: SOCK_*
	bool Create(int domain, int type, int protocol = 0) noexcept
	{
		assert( !IsValid() );
		m_fd = ::socket(domain, type | SOCK_NONBLOCK | SOCK_CLOEXEC, protocol);
		return m_fd != -1; //-1, errno
	}

	//how: SHUT_*
	bool Disable(int how) noexcept
	{
		assert( IsValid() );
		return ::shutdown(m_fd, how) == 0;  //-1, errno
	}

	bool GetSockName(struct sockaddr* addr, socklen_t* addrlen) const noexcept
	{
		assert( IsValid() );
		assert( addr != NULL && addrlen != NULL );
		return ::getsockname(m_fd, addr, addrlen) == 0;  //-1, errno
	}
	bool GetPeerName(struct sockaddr* addr, socklen_t* addrlen) const noexcept
	{
		assert( IsValid() );
		assert( addr != NULL && addrlen != NULL );
		return ::getpeername(m_fd, addr, addrlen) == 0;  //-1, errno
	}

	bool GetOpt(int optname, void* optval, socklen_t* optlen,
				int level = SOL_SOCKET) const noexcept
	{
		assert( IsValid() );
		assert( optval != NULL && optlen != NULL );
		return ::getsockopt(m_fd, level, optname, optval, optlen) == 0;  //-1, errno
	}
	bool SetOpt(int optname, const void* optval, socklen_t optlen,
				int level = SOL_SOCKET) noexcept
	{
		assert( IsValid() );
		assert( optval != NULL && optlen > 0 );
		return ::setsockopt(m_fd, level, optname, optval, optlen) == 0;  //-1, errno
	}

	bool Bind(const struct sockaddr* addr, socklen_t addrlen) noexcept
	{
		assert( IsValid() );
		assert( addr != NULL && addrlen > 0 );
		return ::bind(m_fd, addr, addrlen) == 0;  //-1, errno
	}
	bool Listen(int backlog = SOMAXCONN) noexcept
	{
		assert( IsValid() );
		return ::listen(m_fd, backlog) == 0;  //-1, errno
	}

	//false
	//  code: EAGAIN or EWOULDBLOCK means the operation is still in progress.
	bool Accept(_x_nb_socket& s, int& code, struct sockaddr* addr = NULL, socklen_t* addrlen = NULL) noexcept
	{
		assert( IsValid() );
		assert( !s.IsValid() );
		code = 0;
		int fd = ::accept(m_fd, addr, addrlen);
		if ( fd == -1 ) { //errno
			code = errno;
			return false;
		}
		int ret = 0;
		(void)ret;
		int flags = ::fcntl(fd, F_GETFL);
		if ( flags == -1 ) { //errno
			ret = ::close(fd);
			assert( ret == 0 );  //-1, errno
			return false;
		}
		if ( ::fcntl(fd, F_SETFL, flags | O_NONBLOCK | O_CLOEXEC) == -1 ) { //errno
			ret = ::close(fd);
			assert( ret == 0 );  //-1, errno
			return false;
		}
		s.m_fd = fd;
		return true;
	}
	// EAGAIN means the operation is still in progress on nonblocking UNIX domain sockets.
	// EINPROGRESS means the operation is still in progress.
	bool Connect(const struct sockaddr* addr, socklen_t addrlen) noexcept
	{
		assert( IsValid() );
		assert( addr != NULL && addrlen > 0 );
		return ::connect(m_fd, addr, addrlen) == 0;  //-1, errno
	}

	// recvlen: 0 -- socket peer has performed shutdown
	//   EAGAIN or EWOULDBLOCK means the operation is still in progress.
	bool Recv(void* buf, size_t len, ssize_t& recvlen, int flags = 0) noexcept
	{
		assert( IsValid() );
		recvlen = ::recv(m_fd, buf, len, flags);
		return recvlen != -1;  //-1, errno
	}
	// recvlen: 0 -- socket peer has performed shutdown
	//   EAGAIN or EWOULDBLOCK means the operation is still in progress.
	bool RecvFrom(void* buf, size_t len, ssize_t& recvlen,
			struct sockaddr* src_addr, socklen_t* addrlen,
			int flags = 0) noexcept
	{
		assert( IsValid() );
		assert( src_addr != NULL && addrlen != NULL );
		recvlen = ::recvfrom(m_fd, buf, len, flags, src_addr, addrlen);
		return recvlen != -1;  //-1, errno
	}

	// EAGAIN or EWOULDBLOCK means the operation is still in progress.
	// EPIPE means the peer has closed the connection.
	bool Send(const void* buf, size_t len, ssize_t& sentlen, int flags = MSG_NOSIGNAL) noexcept
	{
		assert( IsValid() );
		sentlen = ::send(m_fd, buf, len, flags);
		return sentlen != -1;  //-1, errno
	}
	// EAGAIN or EWOULDBLOCK means the operation is still in progress.
	// EPIPE means the peer has closed the connection.
	bool SendTo(const void* buf, size_t len, ssize_t& sentlen,
			const struct sockaddr* dest_addr, socklen_t addrlen,
			int flags = MSG_NOSIGNAL) noexcept
	{
		assert( IsValid() );
		assert( dest_addr != NULL && addrlen > 0 );
		sentlen = ::sendto(m_fd, buf, len, flags, dest_addr, addrlen);
		return sentlen != -1;  //-1, errno
	}

protected:
	int m_fd;
};

//------------------------------------------------------------------------------
//tools

inline
bool host_to_sockaddr(const char* host, unsigned int port,
			struct sockaddr_storage& ss, socklen_t& len) noexcept
{
	if ( !_x_socket_helper::FillSockAddr(ss, len, AF_INET, port, host) ) {
		_x_socket_addr sa;
		char szPort[64] = { 0 };
		::snprintf(szPort, sizeof(szPort) / sizeof(char), "%u", port);
		if ( !sa.Fetch(host, szPort) )
			return false;
		const struct addrinfo* pInfo = sa.GetAddrInfo();
		while ( pInfo != NULL ) {
			if ( pInfo->ai_family == AF_INET && pInfo->ai_socktype == SOCK_STREAM ) {
				::memcpy(&ss, pInfo->ai_addr, pInfo->ai_addrlen);
				len = pInfo->ai_addrlen;
				return true;
			}
			pInfo = pInfo->ai_next;
		}
		return false;
	}
	return true;
}

////////////////////////////////////////////////////////////////////////////////
